import networkx as nx
import math

class GraphStats(object):

    def __init__(self, statOp):
        '''
        Constructor
        '''
        self.StatOp = statOp
    
    def Apply(self, network):
        if self.StatOp == "clusteringCoefficient":
            clusteringDict = nx.clustering(nx.Graph(network.graph))
            for node, clusteringCoef in clusteringDict.items():
                node.__dict__["clusteringCoefficient"] = clusteringCoef
                
        elif self.StatOp == "betweenness":
            numNodesToUseInCalculation = max(1, int(math.sqrt(len(network.agentList))))
            betweennessDict = nx.betweenness_centrality(network.graph, numNodesToUseInCalculation)
            
            for node, betweenness in betweennessDict.items():
                node.__dict__["betweenness"] = betweenness
                
        elif self.StatOp == "giantComponentSize":
            largestComponent = nx.connected_component_subgraphs(nx.Graph(network.graph))[0]
            network.__dict__["giantComponentSize"] = largestComponent.number_of_nodes()
            
        elif self.StatOp == "closeness":
            closenessDict = nx.closeness_centrality(network.graph)
            
            for node, closeness in closenessDict.items():
                node.__dict__["closeness"] = closeness
                
        elif self.StatOp == "communicability":
            commDict = nx.communicability_centrality(nx.Graph(network.graph))
            
            for node, comm in commDict.items():
                node.__dict__["communicability"] = comm
        
        elif self.StatOp == "eigenvector":
            graph = self._MakeGraphWtihIntegerNodes(network.graph)
            
            eigenDict = nx.eigenvector_centrality(graph, max_iter=1000, tol=1.0e-4)

            for node, eigenV in eigenDict.items():
                network.agentList[node-1].__dict__["eigenvector"] = eigenV
                
        elif self.StatOp == "currentFlowCloseness":
            #due to a bug in current_flow_closeness_centrality we need to replace all nodes with nodes 
            #of a primitive type.
            graph = self._MakeGraphWtihIntegerNodes(network.graph)
            
            try:
                closenessDict = nx.current_flow_closeness_centrality(graph, weight=None)
                
                for node, closeness in closenessDict.items():
                    network.agentList[node-1].__dict__["currentFlowCloseness"] = closeness

            except Exception:
                for agent in network.agentList:
                    agent.__dict__["currentFlowCloseness"] = -1
            
        elif self.StatOp == "currentFlowBetweenness":
            #due to a bug in current_flow_betweenness_centrality we need to replace all nodes with nodes 
            #of a primitive type.
            graph = self._MakeGraphWtihIntegerNodes(network.graph)
            
            try:
                betweennessDict = nx.current_flow_betweenness_centrality(graph, weight=None)
    
                for node, betweenness in betweennessDict.items():
                    network.agentList[node-1].__dict__["currentFlowBetweenness"] = betweenness
                    
            except Exception:
                for agent in network.agentList:
                    agent.__dict__["currentFlowBetweenness"] = -1
            
        elif self.StatOp == "degree":
            degreeDict = nx.degree(nx.Graph(network.graph))
            
            for node, degree in degreeDict.items():
                node.__dict__["degree"] = degree
            
        else:
            print("Error: '" + str(self.StatOp) + "' is an unknown stat operation.")
            
    def _MakeGraphWtihIntegerNodes(self, sourceGraph):
        graph = nx.Graph(sourceGraph)
        nodeList = graph.nodes()
        nx.relabel_nodes(graph,dict(zip(nodeList,range(1 ,len(nodeList) + 1))), copy=False)
        
        return graph